#! /usr/bin/env python

import os,sys
import argparse
import json
import cv2
from utils.utils import get_yolo_boxes, makedirs
from utils.bbox import get_boxes
from keras.models import load_model
sys.path.append("../")
from xml_factory import XmlMgr
# from tqdm import tqdm
import time

# add by suyongsheng
import tensorflow as tf
from multiprocessing import cpu_count

tf.Session(config=tf.ConfigProto(device_count={"CPU": cpu_count()}, inter_op_parallelism_threads=1,
                                 intra_op_parallelism_threads=2))


def _main_(args):
    config_path = args.conf
    input_path = args.input

    with open(config_path) as config_buffer:
        config = json.load(config_buffer)

    ###############################
    #   Set some parameter
    ###############################       
    net_h, net_w = 416, 416  # a multiple of 32, the smaller the faster
    obj_thresh, nms_thresh = 0.5, 0.45

    ###############################
    #   Load the model
    ###############################
    os.environ['CUDA_VISIBLE_DEVICES'] = config['train']['gpus']
    infer_model = load_model(config['train']['saved_weights_name'])

    image_paths = []
    for dir in os.listdir(input_path):
        sub_dir = os.path.join(input_path, dir)
        if os.path.isdir(sub_dir) == True:
            for file in os.listdir(sub_dir):
                file_path = os.path.join(sub_dir, file)
                extension = os.path.splitext(file_path)[1]
                if extension == ".jpg" \
                    or extension == ".png" \
                    or extension == ".JPEG":
                    image_paths.append(file_path)
    idx=0
    for image_path in image_paths:
        idx+=1
        image = cv2.imread(image_path)
        print(" processing {} of {} at : {}".format(idx , len(image_paths),image_path))

        time_start = time.time()
        # predict the bounding boxes
        boxes = \
            get_yolo_boxes(infer_model, [image], net_h, net_w, config['model']['anchors'], obj_thresh, nms_thresh)[0]
        time_end = time.time()
        print('cost time=', time_end - time_start)

        xyc_pack = get_boxes(image, boxes, config['model']['labels'], obj_thresh)
        xml_gen = XmlMgr(root_node='annotation')
        xml_gen.gen_ann(image_path,xyc_pack,image.shape[1],image.shape[0],image.shape[2])



if __name__ == '__main__':
    argparser = argparse.ArgumentParser(description='Predict with a trained yolo model')
    argparser.add_argument('-c', '--conf', help='path to configuration file')
    argparser.add_argument('-i', '--input', help='path to an image, a directory of images, a video, or webcam')

    args = argparser.parse_args()
    _main_(args)
